package com.example.langchain4j.controller;

import com.example.langchain4j.service.StreamResponse;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.StreamingResponseHandler;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.output.Response;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.Arrays;
import java.util.List;

@RestController
public class OllamaController {

    @Autowired
    private ChatLanguageModel chatModel;

    @Autowired
    private StreamingChatLanguageModel streamingChatModel;

    @Autowired
    private StreamResponse streamResponse;

    @GetMapping("/ai")
    String generation(String userInput) {
        return chatModel.generate(new SystemMessage("请全部用简体中文回答"), new UserMessage(userInput)).content().text();
    }

    @GetMapping(value = "/streamAi", produces = "text/html;charset=UTF-8")
    Flux<String> streamGeneration(@RequestParam(value = "userInput", defaultValue = "给我讲一个笑话") String userInput) {
        return Flux.create(emitter -> {
            List<ChatMessage> chatMessageList = Arrays.asList(new SystemMessage("请全部用简体中文回答"), new UserMessage(userInput));
            streamingChatModel.generate(chatMessageList, new StreamingResponseHandler<>() {
                @Override
                public void onNext(String token) {
                    emitter.next(token); // 发送
                }

                @Override
                public void onComplete(Response<AiMessage> response) {
                    emitter.complete(); // 完成流
                }

                @Override
                public void onError(Throwable error) {
                    emitter.error(error);
                }
            });
        });
    }

}
